package com.kara.woodAgent.agent.tool.provider;

import cn.hutool.core.util.StrUtil;
import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.mcp.client.McpClient;
import dev.langchain4j.service.IllegalConfigurationException;
import dev.langchain4j.service.tool.ToolExecutor;
import dev.langchain4j.service.tool.ToolService;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

import static dev.langchain4j.agent.tool.ToolSpecifications.toolSpecificationFrom;
import static dev.langchain4j.service.IllegalConfigurationException.illegalConfiguration;


@Component
public class ToolProvider implements ApplicationContextAware {

	private final ConcurrentHashMap<String, ToolService> toolServiceMap = new ConcurrentHashMap<>();


	public record ToolServiceHolder(List<ToolService> toolServices, ToolProvider ToolProvider) {


		public List<ToolSpecification> toolSpecifications() {
			return toolServices.stream().map(ToolService::toolSpecifications).flatMap(Collection::stream).toList();
		}

		public List<ToolExecutionResultMessage> executeTool(String memoryId, List<ToolExecutionRequest> requests) {

			return toolServices
					.stream()
					.map(toolService -> {
						Map<String, ToolExecutor> executorMap = toolService.toolExecutors();

						List<ToolExecutionResultMessage> messages = new ArrayList<>();
						for (ToolExecutionRequest request : requests) {
							ToolExecutor executor = executorMap.get(request.name());
							if (executor != null) {
								ToolExecutionResultMessage resultMessage;
								try {
									String execute = executor.execute(request, memoryId);
									if (StrUtil.isBlank(execute)) {
										execute = "执行成功，但是返回内容为空";
									}
									resultMessage = ToolExecutionResultMessage.from(request, execute);
								} catch (Exception e) {
									resultMessage = ToolExecutionResultMessage.from(request, e.getMessage());
									e.printStackTrace();
								}
								messages.add(resultMessage);
							}
						}
						return messages;
					}).flatMap(Collection::stream)
					.collect(Collectors.toList());
		}


	}

	public ToolServiceHolder getToolService(List<String> groups) {

		if (groups != null && !groups.isEmpty()) {
			List<ToolService> toolServices = groups.stream().map(toolServiceMap::get).filter(Objects::nonNull).toList();
			if (toolServices.isEmpty()) {
				return null;
			}
			return new ToolServiceHolder(toolServices, this);

		}

		return null;
	}


	public List<ToolExecutionResultMessage> executeTool(String memoryId, List<String> groups, List<ToolExecutionRequest> requests) {

		return toolServiceMap.entrySet()
				.stream()
				.filter(entry -> groups.contains(entry.getKey()))
				.map(entity -> {
					ToolService toolService = entity.getValue();
					Map<String, ToolExecutor> executorMap = toolService.toolExecutors();

					List<ToolExecutionResultMessage> messages = new ArrayList<>();
					for (ToolExecutionRequest request : requests) {
						ToolExecutor executor = executorMap.get(request.name());
						if (executor != null) {
							ToolExecutionResultMessage resultMessage;
							try {
								resultMessage = ToolExecutionResultMessage.from(request, executor.execute(request, memoryId));
							} catch (Exception e) {
								resultMessage = ToolExecutionResultMessage.from(request, e.getMessage());
								e.printStackTrace();
							}
							messages.add(resultMessage);
						}
					}
					return messages;
				}).flatMap(Collection::stream)
				.collect(Collectors.toList());

	}


	@Override
	public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {

		Map<String, Object> beans = applicationContext.getBeansWithAnnotation(Tools.class);

		Map<String, List<Object>> groups = new HashMap<>();

		for (Map.Entry<String, Object> entity : beans.entrySet()) {

			String key = entity.getKey();
			Object toolObj = entity.getValue();
			Tools annotation = toolObj.getClass().getAnnotation(Tools.class);

			String[] annotationGroups = annotation.groups();

			HashSet<String> groupSet = new HashSet<>(Arrays.asList(annotationGroups));
			groupSet.add(key);
			for (String annotationGroup : groupSet) {
				if (annotationGroup.isBlank()) {
					continue;
				}
				List<Object> objects = groups.computeIfAbsent(annotationGroup, k -> new ArrayList<>());
				objects.add(toolObj);
			}

		}


		//解析
		groups.forEach((group, objects) -> {
			ToolService toolService = new ToolService();
//			toolService.tools(objects);
//			toolService.tools();Map<String,ToolExecutor> toolExecutors = new HashMap<>();

			Map<ToolSpecification, ToolExecutor> toolSpecifications = new HashMap<>();
			for (Object objectWithTool : objects) {
				if (objectWithTool instanceof Class) {
					throw illegalConfiguration("Tool '%s' must be an object, not a class", objectWithTool);
				}

				for (Method method : objectWithTool.getClass().getDeclaredMethods()) {
					if (method.isAnnotationPresent(Tool.class)) {
						ToolSpecification toolSpecification = toolSpecificationFrom(method);
						if (toolSpecifications.keySet().stream().map(ToolSpecification::name).anyMatch(a -> a.equals(toolSpecification.name()))) {
							throw new IllegalConfigurationException(
									"Duplicated definition for tool: " + toolSpecification.name());
						}
						toolSpecifications.put(toolSpecification, new ConsumeToolExecutor(objectWithTool, method));
					}
				}
			}
			toolService.tools(toolSpecifications);
			//解析mcp客户端
			parseMcpClient(toolService, objects);
			toolServiceMap.put(group, toolService);
		});

	}

	private void parseMcpClient(ToolService toolService, List<Object> objects) {

		for (Object object : objects) {
			for (Method method : object.getClass().getSuperclass().getDeclaredMethods()) {
				McpTool annotation = method.getAnnotation(McpTool.class);
				if (annotation != null && annotation.enable()) {
					try {
						McpClient mcpClient = (McpClient) method.invoke(object);
						Map<ToolSpecification, ToolExecutor> executorMap = getExecutorMap(mcpClient);
						toolService.tools(executorMap);
					} catch (IllegalAccessException | InvocationTargetException e) {
						throw new RuntimeException(e);
					}

				}
			}
		}

	}

	private final Map<McpClient, Map<ToolSpecification, ToolExecutor>> executorCache = new HashMap<>();

	@NotNull
	private Map<ToolSpecification, ToolExecutor> getExecutorMap(McpClient mcpClient) {

		if (executorCache.containsKey(mcpClient)) {
			return executorCache.get(mcpClient);
		}

		List<ToolSpecification> toolSpecifications = mcpClient.listTools();
		Map<ToolSpecification, ToolExecutor> executorMap = toolSpecifications.stream()
				.map(toolSpecification -> new McpToolExecutor(mcpClient, toolSpecification))
				.collect(Collectors.toMap(McpToolExecutor::toolSpecification, mcpToolExecutor -> mcpToolExecutor,
						(oldValue, newValue) -> oldValue));

		executorCache.put(mcpClient, executorMap);
		return executorMap;
	}
}
